import os
from ACGAN import ACGAN
from tools import checkFolder
import tensorflow as tf
import argparse
import numpy as np


def main(args):

    if args is None :
        print("命令行参数错误！")
        exit(0)
    """
        def __init__(self,sess,epoch,batchSize,codeSize,dataSource,\
        checkpointDir,resultDir,logDir,mode,height,width,Nchannel):
    """
    with tf.Session() as sess :
        myGAN = ACGAN(sess,args.epoch,args.batchSize,args.codeSize,\
            args.dataSource,args.checkpointDir,args.resultDir,args.logDir,args.mode,\
                64,64,3)
        if myGAN is None:
            print("创建GAN网络失败")
            exit(0)
        
        if args.mode=='train' :
            myGAN.buildNet()
            print("进入训练模式")
            myGAN.train()
            print("Done")
        elif args.mode=='infer' :
            myGAN.buildForInfer()
            tag_dict = ['orange hair', 'white hair', 'aqua hair', 'gray hair', 'green hair', 'red hair', 'purple hair', 'pink hair', 'blue hair', 'black hair', 
                        'brown hair', 'blonde hair','gray eyes', 'black eyes', 'orange eyes', 'pink eyes', 'yellow eyes','aqua eyes', 'purple eyes', 'green eyes',
                        'brown eyes', 'red eyes','blue eyes']
            tag = np.zeros((64,23))
            feature = args.hairStyle+" AND "+ args.eyeStyle
            for j in range(25):
                for i in range(len(tag_dict)):
                    if tag_dict[i] in feature:
                        tag[j][i] = 1
            myGAN.infer(tag,feature)
            print("Generate : "+feature)

